/*
Dwarf Therapist
Copyright (c) 2009 Trey Stout (chmod)

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
*/
#include <QtGui>
#include <QtDebug>

#include "dfinstance.h"
#include "dfinstanceosx.h"
#include "defines.h"
#include "dwarf.h"
#include "utils.h"
#include "gamedatareader.h"
#include "memorylayout.h"
#include "cp437codec.h"
#include "memorysegment.h"
#include "truncatingfilelogger.h"
#include <stdio.h>
#include <mach/vm_map.h>
#include <mach/mach_traps.h>
#include <mach-o/dyld.h>
#include <mach/task.h>
#include <SecurityFoundation/SFAuthorization.h>
#include <Security/AuthorizationTags.h>
#include <CommonCrypto/CommonDigest.h>
#include <openssl/md5.h>

DFInstanceOSX::DFInstanceOSX(QObject* parent)
    : DFInstance(parent)    
{ 
}

DFInstanceOSX::~DFInstanceOSX() {
}

uint DFInstanceOSX::calculate_checksum() {
    // pool is created
    NSAutoreleasePool *authPool = [[NSAutoreleasePool alloc] init]; 
    
    quint32 md5 = 0xe3775089;

    QString str_md5;

    unsigned char outputData[CC_MD5_DIGEST_LENGTH];
    
    NSData *inputData = [[NSData alloc] initWithContentsOfFile:
                            [NSString stringWithCString:m_loc_of_dfexe.toUtf8()
                            length:strlen( m_loc_of_dfexe.toUtf8()) ]];
    
    CC_MD5([inputData bytes], [inputData length], outputData);
    [inputData release];
    
    NSMutableString *hash = [[NSMutableString alloc] init];
    
    for (NSUInteger i = 0; i < CC_MD5_DIGEST_LENGTH; i++) {
        [hash appendFormat:@"%02x", outputData[i]];
    }
    
    str_md5 = QString([hash UTF8String]);
    bool ok;
    md5 = str_md5.mid(0, 8).toUInt(&ok,16); // use the first 8 bytes

    TRACE << "GOT MD5:" << str_md5;
    
    [authPool release];
    
    return md5;
}

QVector<VIRTADDR> DFInstanceOSX::enumerate_vector(const VIRTADDR &addr) {
    QVector<uint> addrs;
    if (!addr)
        return addrs;
    
    VIRTADDR start = read_addr(addr);
    VIRTADDR end = read_addr(addr + 4);
    int bytes = end - start;
    int entries = bytes / 4;
    TRACE << "enumerating vector at" << hex << addr << "START" << start
    << "END" << end << "UNVERIFIED ENTRIES" << dec << entries;
    VIRTADDR tmp_addr = 0;
    
    if (entries > 5000) {
        LOGW << "vector at" << hexify(addr) << "has over 5000 entries! (" <<
        entries << ")";
    }
    
    QByteArray data(bytes, 0);
    int bytes_read = read_raw(start, bytes, data);
    if (bytes_read != bytes && m_layout->is_complete()) {
        TRACE << "Tried to read" << bytes << "bytes but only got"
        << bytes_read;
        return addrs;
    }
    for(int i = 0; i < bytes; i += 4) {
        tmp_addr = decode_dword(data.mid(i, 4));
        if (m_layout->is_complete()) {
            if (is_valid_address(tmp_addr))
                addrs << tmp_addr;
        } else {
            addrs << tmp_addr;
        }
    }
    detach();
    return addrs;
}

QString DFInstanceOSX::read_string(const uint &addr) {
    VIRTADDR buffer_addr = read_addr(addr);
    int upper_size = 256;
    QByteArray buf(upper_size, 0);
    read_raw(buffer_addr, upper_size, buf);
    
    QString ret_val(buf);
    CP437Codec *codec = new CP437Codec;
    ret_val = codec->toUnicode(ret_val.toAscii());
    return ret_val;
}

int DFInstanceOSX::write_string(const VIRTADDR &addr, const QString &str) {
    Q_UNUSED(addr);
    Q_UNUSED(str);
    return 0;
}

int DFInstanceOSX::write_int(const VIRTADDR &addr, const int &val) {
    return write_raw(addr, sizeof(int), (void*)&val);
}

int DFInstanceOSX::read_raw(const VIRTADDR &addr,
                            int bytes, 
                            QByteArray &buffer) {

    kern_return_t result;
        
    vm_size_t readsize = 0;
        
    result = vm_read_overwrite(m_task, addr, bytes,
                               (vm_address_t)buffer.data(),
                               &readsize );
    
    if ( result != KERN_SUCCESS ) {
        return 0;
    }
    
    return readsize;
    
}

int DFInstanceOSX::write_raw(const VIRTADDR &addr, const int &bytes, 
                             void *buffer) {
    kern_return_t result;

    result = vm_write( m_task,  addr,  (vm_offset_t)buffer,  bytes );
    if ( result != KERN_SUCCESS ) {
        return 0;
    }

    return bytes;
}

bool DFInstanceOSX::find_running_copy() {
    NSAutoreleasePool *authPool = [[NSAutoreleasePool alloc] init];  
    
    NSWorkspace *workspace = [NSWorkspace sharedWorkspace];
    NSArray *launchedApps = [workspace launchedApplications];
    unsigned i, len = [launchedApps count];
        
    // compile process array
    for ( i = 0; i < len; i++ ) {
        NSDictionary *application = [launchedApps objectAtIndex:i];
        if ( [[application objectForKey:@"NSApplicationName"] 
                                            isEqualToString:@"dwarfort.exe" ]) {
            
            m_loc_of_dfexe = QString( [[application objectForKey:
                                     @"NSApplicationPath"] UTF8String]);
            m_pid = [[application objectForKey:
                                            @"NSApplicationProcessIdentifier"] 
                                            intValue];
        }
    }
    TRACE << "Location of exe" << m_loc_of_dfexe;
    if (m_pid == 0 || ! (task_for_pid( current_task(), m_pid, &m_task ) 
                                                            == KERN_SUCCESS)) {
        QMessageBox::warning(0, tr("Warning"),
                             tr("Unable to locate a running copy of Dwarf "
                                "Fortress, are you sure it's running?"));
        LOGW << "can't find running copy";
        m_is_ok = false;
        return m_is_ok;
        
    }
    m_layout = get_memory_layout(hexify(calculate_checksum()).toLower());
    //LOGD << "PEB is at: " << hex << peb_addr;
    m_is_ok = true;
    map_virtual_memory();
    
    [authPool release];
    
    return true;
}

BYTE DFInstanceOSX::read_byte(const VIRTADDR &addr) {
    QByteArray out;
    read_raw(addr, sizeof(BYTE), out);
    return (BYTE) out.data()[0];
}

WORD DFInstanceOSX::read_word(const VIRTADDR &addr) {
    QByteArray out;
    read_raw(addr, sizeof(WORD), out);
    return decode_word(out);
}

/*! OS specific way of asking the kernel for valid virtual memory pages from
 the DF process. These pages are used to speed up scanning, and validate
 reads from DF's memory. If addresses are not within ranges found by this
 method, they will fail the is_valid_address() method */
void DFInstanceOSX::map_virtual_memory() {    
    if (m_regions.isEmpty()) {        
        if (!m_is_ok)
            return;
        
        m_lowest_address = 0x1000;
        m_highest_address = 0xFFFFFFFF;
        MemorySegment *segment = new MemorySegment(QString("[stack]"),
                                            m_lowest_address, m_highest_address);
        m_regions << segment;
        LOGW << segment->to_string();
    }
}


bool DFInstanceOSX::attach() {
    kern_return_t result;
    
    result = task_suspend(m_task);
    
    if ( result != KERN_SUCCESS ) {
        return false;
    }
    
    return true;
    
}

bool DFInstanceOSX::detach() {
    kern_return_t result;
    result = task_resume(m_task);
    if ( result != KERN_SUCCESS ) {
        return false;
    }
    return true; 
}


/* This will allow us to read a write memory and root.
   Will bring up standard OSX
   SecurtyFramework dialog */
void DFInstanceOSX::authMe()
{
    
    NSAutoreleasePool *authPool = [[NSAutoreleasePool alloc] init]; 
    
    NSString *myPath = [[NSBundle mainBundle] resourcePath];
    chdir([myPath UTF8String]);   
        
    AuthorizationRef myAuthRef;
    OSStatus stat = AuthorizationCopyPrivilegedReference(&myAuthRef,
                                                    kAuthorizationFlagDefaults);
    NSDictionary    *applicationAttributes = [[NSFileManager defaultManager] 
                                              fileAttributesAtPath:
                                              [[NSBundle mainBundle] 
                                              executablePath] 
                                              traverseLink: YES];
    
    if(!( stat == errAuthorizationSuccess || 
                ([applicationAttributes filePosixPermissions] == 1517 && 
                 [[applicationAttributes fileGroupOwnerAccountName]
                                            isEqualToString: @"procmod"]))) {
        // get authorization as root
        
        char FullPathToMe[1024];
        uint32_t size = sizeof(FullPathToMe);
        if (_NSGetExecutablePath(FullPathToMe, &size) == 0)
            printf("executable path is %s\n", FullPathToMe);
        else
            printf("buffer too small; need size %u\n", size);    
        
        OSStatus myStatus;
        
        // set up Authorization Item
        AuthorizationItem myItems[1];
        myItems[0].name = kAuthorizationRightExecute;
        myItems[0].valueLength = 0;
        myItems[0].value = NULL;
        myItems[0].flags = 0;
        
        // Set up Authorization Rights
        AuthorizationRights myRights;
        myRights.count = sizeof (myItems) / sizeof (myItems[0]);
        myRights.items = myItems;
        
        // set up Authorization Flags
        AuthorizationFlags myFlags;
        myFlags =
        kAuthorizationFlagDefaults |
        kAuthorizationFlagInteractionAllowed |
        kAuthorizationFlagExtendRights;
        
        // Create an Authorization Ref using Objects above. 
        // NOTE: Login bod comes up with this call.
        AuthorizationRef myAuthorizationRef;
        myStatus = AuthorizationCreate (&myRights,
                                        kAuthorizationEmptyEnvironment,
                                        myFlags, &myAuthorizationRef);
        
        if (myStatus == errAuthorizationSuccess)
        {
            // prepare communication path
            // used to signal that process is loaded
            FILE *myCommunicationsPipe = NULL;
            char myReadBuffer[] = " ";
            
            myStatus = AuthorizationExecuteWithPrivileges(myAuthorizationRef,
                                                    FullPathToMe,
                                                    kAuthorizationFlagDefaults,
                                                    nil,
                                                    &myCommunicationsPipe);
            
            // external app is running asynchronously
            // - it will send to stdout when loaded
            if (myStatus == errAuthorizationSuccess)
            {
                read (fileno (myCommunicationsPipe), myReadBuffer, 
                                                    sizeof (myReadBuffer));
                fclose(myCommunicationsPipe);
            }
            
            // release authorization reference
            myStatus = AuthorizationFree (myAuthorizationRef, 
                                          kAuthorizationFlagDestroyRights);
        }
        [authPool release];
        exit(0);
    }
    
        
}
